/**
  ******************************************************************************
  * @file           : mem_leak2.cpp
  * @author         : JarryLi
  * @brief          : None
  * @date           : 2024/10/14 9:24
  ******************************************************************************
  */

#include <dlfcn.h>
#include <link.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#define BUFF_SIZE 128

#if 1

// 函数指针
typedef void *(*malloc_t)(size_t size);
// 定义函数
malloc_t malloc_f = NULL;

// 函数指针
typedef void (*free_t)(void *ptr);
// 定义函数
free_t free_f = NULL;

int enable_malloc = 1;
int enable_free = 1;

void* ConvertToELF(void* addr)
{
    // 定义一个Dl_info结构体变量，用于存储与地址关联的动态链接信息
    Dl_info info;

    // 定义一个link_map结构体指针，用于存储动态链接映射信息
    struct link_map* link;

    // 调用dladdr1函数，获取地址addr相关的动态链接信息。
    // addr是我们要查询的地址。
    // &info是Dl_info结构体的指针，用于存储查询结果。
    // (void**)&link是一个指向void指针的指针，用于接收link_map结构体的指针。
    // RTLD_DL_LINKMAP是一个标志，告诉dladdr1函数我们需要link_map结构体。
    dladdr1(addr, &info, (void**)&link, RTLD_DL_LINKMAP);

    // 计算并返回转换后的地址。这里将原始地址addr减去link_map结构体中的l_addr成员，
    // l_addr是映射的基地址，这样可以得到相对于ELF文件开始的偏移量。
    return (void*)((size_t)addr - link->l_addr);
}

void* malloc(size_t size)
{
    void* p = NULL;
    // 内部会调用 malloc 和 free
    if (enable_malloc)
    {
        enable_malloc = 0;
        p = malloc_f(size);

        // 获取当前函数调用者的地址
        //      参数为0: 调用者(上一个)
        //      参数为1: 上上个调用者
        void* caller = __builtin_return_address(0);

        char buff[BUFF_SIZE] =  { 0 };
        sprintf(buff, "./file/%p.mem", p);

        FILE* fp = fopen(buff, "w");
        if (!fp)
        {
            free(p);
            return NULL;
        }

        printf("[+]%p, addr: %p, size: %ld\n", ConvertToELF(caller), p, size);
        fprintf(fp, "[+]%p, addr: %p, size: %ld\n", ConvertToELF(caller), p, size);
        fflush(fp);

        enable_malloc = 1;
    }
    else
    {
        p = malloc_f(size);
    }

    return p;
}

void free(void* ptr)
{
    // 内部会调用 malloc 和 free
    if (enable_free)
    {
        enable_free = 0;
        char buff[BUFF_SIZE] = { 0 };
        snprintf(buff, BUFF_SIZE, "./file/%p.mem", ptr);

        if (unlink(buff) < 0)
        {
            printf("double free: %p", ptr);
            return;
        }

        free_f(ptr);

        enable_free = 1;
    }
    else
    {
        free_f(ptr);
    }

    return;
}

void init_hook()
{
    if (!malloc_f)
    {
        // 如果malloc_f没有被初始化，使用dlsym函数来查找并获取下一个（即原始的）malloc函数的地址。
        // RTLD_NEXT是一个特殊的句柄，它告诉dlsym函数在当前程序的链接顺序中查找下一个可用的符号。
        // "malloc"是我们要查找的符号名，dlsym返回这个符号的地址。
        // (malloc_t)dlsym(...)是一个类型转换，将返回的函数指针转换为malloc_t类型的指针。
        malloc_f = (malloc_t)dlsym(RTLD_NEXT, "malloc");
    }

    if (!free_f)
    {
        free_f = (free_t)dlsym(RTLD_NEXT, "free");
    }
}

#endif

int main() {

    init_hook();

#if 1
    void *p1 = malloc(5);
    void *p2 = malloc(10);
    void *p3 = malloc(35);
    void *p4 = malloc(10);

    free(p1);
    free(p3);
    free(p4);
#endif
    getchar(); // 防止死递归

}
